import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet50, ResNet50_Weights


class GlobalAveragePooling(nn.Module):
    """全局平均池化模块"""

    def forward(self, x):
        return nn.AdaptiveAvgPool2d(output_size=(1, 1))(x)


class ParseNet(nn.Module):
    """ParseNet 模型"""

    def __init__(self, num_classes):
        super(ParseNet, self).__init__()
        # 使用 ResNet50 作为骨干网络
        base_model = resnet50(weights=ResNet50_Weights.DEFAULT)
        self.backbone = nn.Sequential(*list(base_model.children())[:-2])  # 去掉最后两层

        # 全局上下文特征模块
        self.global_context = nn.Sequential(
            GlobalAveragePooling(),
            nn.Conv2d(2048, 512, kernel_size=1, stride=1, bias=False),
            # nn.BatchNorm2d(512),
            nn.ReLU()
        )

        # 解码模块
        self.decoder = nn.Sequential(
            nn.Conv2d(2560, 512, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(),
            nn.Conv2d(512, num_classes, kernel_size=1)
        )

    def forward(self, x):
        # 主干网络提取特征
        features = self.backbone(x)  # 输出大小: [B, 2048, H/32, W/32]

        # 提取全局上下文特征
        global_context = self.global_context(features)  # [B, 512, 1, 1]
        global_context = F.interpolate(global_context, size=features.size()[2:], mode='bilinear', align_corners=True)

        # 将全局特征和局部特征拼接
        combined_features = torch.cat([features, global_context], dim=1)  # [B, 2560, H/32, W/32]

        # 解码为最终的分割结果
        out = self.decoder(combined_features)  # [B, num_classes, H/32, W/32]
        out = F.interpolate(out, size=x.size()[2:], mode='bilinear', align_corners=True)  # 恢复原图大小

        return out


if __name__ == '__main__':
    x = torch.randn(1, 3, 224, 224)
    model = ParseNet(num_classes=10)
    y = model(x)
    print(y.size())
